#!/usr/bin/env python
# -*- coding=utf-8 -*-

#This script can merge multiple vcfs in multiple process by SGE
import sys
import os
import argparse
import subprocess
from uuid import uuid1
from configparser import ConfigParser
from textwrap import dedent
from glob import glob

current_dir = os.path.dirname(os.path.abspath(__file__))

def make_tmp_tree(tmp_dir):
    """
    make tmp tree directories to save middle files and scripts
    tmp tree:
    |--tmp_node
        |--scripts
        |--tmpfiles
        |--intervals
        |--job
            |--log
    """
    tmp_name = str(uuid1())
    tmp_node = os.path.join(tmp_dir,tmp_name)
    scripts_node = os.path.join(tmp_node,'scripts')
    midfile_node = os.path.join(tmp_node,'tmpfiles')
    intervals_node = os.path.join(tmp_node,'intervals')
    job_node = os.path.join(tmp_node,'job')
    os.makedirs(tmp_node)
    os.mkdir(scripts_node)
    os.mkdir(midfile_node)
    os.mkdir(intervals_node)
    os.mkdir(job_node)
    os.mkdir(os.path.join(job_node,'log'))

    return tmp_node,scripts_node,midfile_node,intervals_node,job_node

    
def get_vcf_list(vcf_list_file):
    """
    get vcf files from vcf_list file
    """
    vcf_list = []
    with open(vcf_list_file,'r') as indata:
        for line in indata:
            vcf_list.append(line.strip())
    return vcf_list

def make_merge_script(merge_software,vcf_list_file,interval_file,midfile_node,scripts_node,other_args,all_configs):
    """
    merge_software : gatk or bcftools
    """
    ofile = os.path.join(midfile_node,os.path.basename(interval_file) + '.vcf.gz')
    oshell = os.path.join(scripts_node,"work."+ os.path.basename(interval_file) + '.merge.sh')
    if merge_software == 'bcftools':
        script = dedent("""
        set -eo pipefail
        echo start `date +"%F %T"`
        {bcftools} merge\\
            -l {vcf_list_file} \\
            -g {ref} \\
            -R {interval_file} \\
            -Oz -o {ofile} \\
            {other_args}
        echo end `date +"%F %T"`
        """.format(vcf_list_file=vcf_list_file,interval_file=interval_file,
            ofile=ofile,other_args=other_args,**all_configs))
    elif merge_software == 'gatk-CombineGVCFs':
        vcfs_list = get_vcf_list(vcf_list_file)
        variants_args = "".join(["--variant {} ".format(vcf) for vcf in vcfs_list])
        script = dedent("""
        set -eo pipefail
        echo start `date +"%F %T"`
        {gatk} CombineGVCFs \\
            -R {ref} \\
            {variants_args} \\
            -L {interval_file} \\
            -O {ofile} \\
            {other_args}
        echo end `date +"%F %T"`
        """.format(variants_args=variants_args,interval_file=interval_file,
            ofile=ofile,other_args=other_args,**all_configs))
    elif merge_software == "gatk-GenomicsDBImport":
        script = dedent("""
        set -eo pipefail
        echo start `date +"%F %T"`
        {gatk} GenomicsDBImport \\
            --sample-name-map {vcf_list_file} \\
            --genomicsdb-workspace-path {midfile_node}/my_database \\
            --tmp-dir {midfile_node}/tmp \\
            -L {interval_file}
        echo end `date +"%F %T"`
        """.format(vcf_list_file=vcf_list_file,midfile_node=midfile_node,
            interval_file=interval_file))
    else:
        sys.stderr.write('Unknown merge software {merge_software},please check!'.format(
            merge_software=merge_software))
        exit(1)

    with open(oshell,'w') as odata:
        odata.write(script)
    print('output merge script to {oshell}'.format(oshell=oshell))

    return oshell,ofile

def concat_interval_vcfs(ofile_list,scripts_node,merge_vcf_ofile,all_configs):
    """
    Concatenate the interval vcfs generated by bcftools merge 
    """
    oshell = os.path.join(scripts_node,'work.concat.sh')
    script = dedent("""
    {bcftools} concat \\
        -Oz -o {merge_vcf_ofile} \\
        {ofile_list}
    """.format(
        merge_vcf_ofile=merge_vcf_ofile,
        bcftools = all_configs.get('bcftools'),
        ofile_list= " ".join(ofile_list)
    ))
    with open(oshell,'w') as odata:
        odata.write(script)
    print("output concatenating vcfs work shell to {oshell}".format(oshell=oshell))

    return oshell
    
def gatk_genotypegvcfs(merge_software,midfile_node,scripts_node,merge_vcf_ofile,all_configs,split_vcfs_list):
    """
    in the gatk pipeline files generated by CombineGVCFs and GenomicsDBImport
    commands need to perform joint genotyping 
    """
    oshell = os.path.join(scripts_node,"work.genotypegvcfs.sh")
    if merge_software == "gatk-CombineGVCFs":
        script = dedent("""
        {gatk} GatherVcfs \\
            {split_vcfs} \\
            -O {midfile_node}/gather.g.vcf.gz

        {bcftools} index -t {midfile_node}/gather.g.vcf.gz
        
        {gatk} GenotypeGVCFs \\
            -R {ref} \\
            -V {midfile_node}/gather.g.vcf.gz \\
            -O {merge_vcf_ofile} \\

        """.format(split_vcfs=" ".join(["-I {}".format(v) for v in split_vcfs_list]),
        midfile_node=midfile_node,merge_vcf_ofile=merge_vcf_ofile,**all_configs))
    elif merge_software == "gatk-GenomicsDBImport":
        script = dedent("""
        {gatk} GenotypeGVCFs \\
            -R {ref} \\
            -V gendb://{midfile_node}/my_database/ \\
            -O {merge_vcf_ofile} \\
            --tmp-dir {midfile_node}/tmp
        """.format(midfile_node=midfile_node,merge_vcf_ofile=merge_vcf_ofile,
            scripts_node=scripts_node,**all_configs))

    with open(oshell,'w') as odata:
        odata.write(script)
    print("output genotypegvcf work shell to {oshell}".format(oshell=oshell))

    return oshell

def split_intervals(scatter_number,scripts_node,intervals_node,all_configs):
    """
    split reference by given numbers
    """
    oshell = os.path.join(scripts_node,"work.split_interval.sh")
    script = dedent("""
    #make intervals for gatk
    {gatk} SplitIntervals \\
        -R {ref} \\
        -L {ref_range} \\
        --scatter-count {scatter_number} \\
        -O {intervals_node}

    #make simple interval for bcftools
    ls {intervals_node}/*|xargs -i echo "grep -v '@' {{}} > {{}}.simple"|sh
    """.format(scatter_number=scatter_number,
        intervals_node=intervals_node,**all_configs))

    with open(oshell,'w') as odata:
        odata.write(script)
    print('output split interval script to {oshell}'.format(oshell=oshell))

    return oshell


def makejob(job_list,job_node,all_configs):
    """
    generate job config and job file
    """
    oconfig = os.path.join(job_node,'job_config.txt')
    outfile = os.path.join(job_node,'splitBigData.job')
    queue = all_configs.get('submit_queues')
    logdir = os.path.join(job_node,'log')
    with open(oconfig,'w') as odata:
        for line in job_list:
            odata.write("\t".join(line) + "\n")
    mkjob_cmd = "{current_dir}/makejob.py {oconfig} --outfile {outfile} --queue {queue} --logdir {logdir}".format(
        oconfig=oconfig,current_dir = current_dir,outfile = outfile,queue = queue,logdir = logdir
    )
    # print(mkjob_cmd)
    assert not subprocess.call(mkjob_cmd,shell=True)
    # print('Generate job to {outfile}'.format(outfile=outfile))
    # assert not subprocess.call('sjm {outfile}'.format(outfile=outfile),shell=True)

def sort_interval(key):
    return int(os.path.basename(key).split('-')[0])

def merge_vcf_main(args):
    """
    the pipeline of merging vcf
    """
    config_file = args.get('config')
    tmp_dir = os.path.abspath(args.get('tmp_dir'))
    scatter_number = args.get('scatter_number')
    merge_software = args.get('merge_software')
    vcfs_list_file = os.path.abspath(args.get('vcfs_list'))
    print(vcfs_list_file)
    other_args = args.get('other_args')
    merge_vcf_ofile = os.path.abspath(args.get('ofile'))

    config = ConfigParser()
    if not config:
        config.read(os.path.join(current_dir,'config.ini'))
    else:
        config.read(config_file)

    softwares = dict(config.items("software"))
    reference = dict(config.items("reference"))
    others = dict(config.items("others"))

    all_configs = {}
    for section in [softwares,reference,others]:
        all_configs.update(section)

    job_list = []
    split_vcfs_script_list = []
    split_vcfs_list = []
    tmp_node,scripts_node,midfile_node,intervals_node,job_node = make_tmp_tree(tmp_dir)
    split_oshell = split_intervals(scatter_number,scripts_node,intervals_node,all_configs)

    assert not subprocess.call("sh {split_oshell}".format(split_oshell=split_oshell),shell=True)
    # job_list.append([split_oshell,'2G',])
    ## split intervals
    # interval_files = [os.path.join(intervals_node,i) for i in os.listdir(intervals_node)]
    if merge_software == "bcftools":
        interval_files = glob("{intervals_node}/*.simple".format(intervals_node=intervals_node))
    else:
        interval_files = glob("{intervals_node}/*.interval_list".format(intervals_node=intervals_node))

    interval_files.sort(key=sort_interval)
    ## merge vcfs by split intervals
    for interval_file in interval_files:

        oshell,ofile = make_merge_script(merge_software,vcfs_list_file,interval_file,midfile_node,scripts_node,other_args,all_configs)
        job_list.append([oshell,all_configs.get('merge_memory','5G')])
        split_vcfs_list.append(ofile)
        split_vcfs_script_list.append(oshell)
    
    ## concate vcfs
    if merge_software == "bcftools":
        concat_oshell = concat_interval_vcfs(split_vcfs_list,scripts_node,merge_vcf_ofile,all_configs)
        job_list.append([concat_oshell,all_configs.get('concat_memory','2G'),",".join(split_vcfs_script_list),])
    if merge_software.startswith('gatk-'):
        oshell = gatk_genotypegvcfs(merge_software,midfile_node,scripts_node,merge_vcf_ofile,all_configs,
            split_vcfs_list)
        job_list.append([oshell,all_configs.get('concat_memory','2G'),",".join(split_vcfs_script_list)])
    
    ## make job
    # print(job_list)
    makejob(job_list,job_node,all_configs)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--config","-c",help="basic configurations about merge vcf")
    parser.add_argument("--tmp_dir","-t",default=os.environ["PWD"],
        help="where to generate temporary directory default is current directory")
    parser.add_argument("--ofile",'-o',help='filename of output merged vcfs')
    parser.add_argument("--vcfs_list",'-v',help="a file contains vcfs' path that prepared to merge quickly")
    parser.add_argument("--scatter_number",'-s',default=30,help="how many intervals to split to merge parallel[default:30]")
    parser.add_argument("--merge_software",'-m',default="bcftools",choices=['bcftools','gatk-CombineGVCFs','gatk-GenomicsDBImport'],
        help="software applied to merge vcf,choose from ['bcftools','gatk-CombineGVCFs','gatk-GenomicsDBImport']")
    parser.add_argument("--other_args",'-args',default="",help="other arguments you want to pass to merge software")


    args = vars(parser.parse_args())
    merge_vcf_main(args)